package cn.chendd.cofig.filters;

import cn.chendd.cofig.enums.HttpParamScopeEnum;
import io.netty.buffer.ByteBufAllocator;
import org.apache.commons.lang3.StringUtils;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.core.io.buffer.*;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerCodecConfigurer;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.annotation.Resource;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;

/**
 * 缓存body 过滤器
 * @author chendd
 * @date 2022/12/30 17:14
 */
@Component
public class GatewayContextFilter implements GlobalFilter, Ordered {

    @Resource
    private ServerCodecConfigurer codecConfigurer;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        MediaType mediaType = exchange.getRequest().getHeaders().getContentType();
        if (mediaType == null || MediaType.APPLICATION_FORM_URLENCODED.isCompatibleWith(mediaType)) {
            GatewayContext gatewayContext = new GatewayContext();
            exchange.getAttributes().put(GatewayContext.REQUEST_BODY_CACHE_GATEWAY_CONTEXT, gatewayContext);
            return readFormData(exchange, chain, gatewayContext);
        }
        HttpHeaders headers = exchange.getRequest().getHeaders();
        MediaType contentType = headers.getContentType();
        long contentLength = headers.getContentLength();
        //处理包含请求体的逻辑
        if (contentLength > 0) {
            //处理请求
            GatewayContext gatewayContext = new GatewayContext();
            exchange.getAttributes().put(GatewayContext.REQUEST_BODY_CACHE_GATEWAY_CONTEXT, gatewayContext);
            if (MediaType.APPLICATION_JSON.isCompatibleWith(contentType)) {
                return readBody(exchange, chain, gatewayContext);
            }
            //若为文件上传
            if (MediaType.MULTIPART_FORM_DATA.isCompatibleWith(contentType)) {
                return readBody(exchange, chain, gatewayContext);
            }
        }
        //处理json无请求体的请求
        if (MediaType.APPLICATION_JSON.isCompatibleWith(contentType)) {
            GatewayContext gatewayContext = new GatewayContext();
            exchange.getAttributes().put(GatewayContext.REQUEST_BODY_CACHE_GATEWAY_CONTEXT, gatewayContext);
            return readFormData(exchange, chain, gatewayContext);
        }
        return this.doChainFilter(exchange , chain , null);
    }

    @Override
    public int getOrder() {
        return -100;
    }

    /**
     * 过滤器放行
     * @param exchange 交换机
     * @param chain 过滤器链
     * @param gatewayContext 数据对象
     * @return 数据对象
     */
    private Mono<Void> doChainFilter(ServerWebExchange exchange , GatewayFilterChain chain , GatewayContext gatewayContext) {
        if (gatewayContext == null) {
            return chain.filter(exchange);
        }
        return chain.filter(exchange).doOnSuccess(item -> {
            disposeRequestData(exchange);
        });
    }

    /**
     * 处理请求的数据日志
     * @param exchange 交换机
     */
    private void disposeRequestData(ServerWebExchange exchange) {
        List<String> httpParamScope = exchange.getAttribute(GatewayContext.HTTP_PARAM_SCOPE);
        HttpParamScopeEnum.execute(httpParamScope , exchange);
    }

    /**
     * 读取form data 缓存至body
     * @param exchange 交换机
     * @param chain 过滤器链
     * @param gatewayContext 数据对象
     * @return 处理结果
     */
    private Mono<Void> readFormData(ServerWebExchange exchange, GatewayFilterChain chain, GatewayContext gatewayContext) {
        HttpHeaders headers = exchange.getRequest().getHeaders();
        return exchange.getFormData()
                .doOnNext(gatewayContext::setFormData)
                .then(Mono.defer(() -> {
                    MediaType mediaType = headers.getContentType();
                    Charset charset = StandardCharsets.UTF_8;
                    if (mediaType != null) {
                        Charset c = mediaType.getCharset();
                        if (c != null) {
                            charset = c;
                        }
                    }
                    MultiValueMap<String, String> formData = gatewayContext.getFormData();
                    //当不存在参数值时，返回
                    if (formData == null || formData.isEmpty()) {
                        ServerHttpResponseDecorator decoratedResponse = this.getServerHttpResponseDecorator(exchange , gatewayContext);
                        return chain.filter(exchange.mutate().response(decoratedResponse).build()).doOnSuccess(item -> {
                            disposeRequestData(exchange);
                        });
                    }
                    StringBuilder formDataBodyBuilder = new StringBuilder();
                    try {
                        String entryKey;
                        List<String> entryValue;
                        for (Map.Entry<String, List<String>> entry : formData.entrySet()) {
                            entryKey = entry.getKey();
                            entryValue = entry.getValue();
                            if (entryValue.size() > 1) {
                                for (String value : entryValue) {
                                    formDataBodyBuilder.append(entryKey).append("=").append(URLEncoder.encode(value, charset.name())).append("&");
                                }
                            } else if (entryValue.size() > 0) {
                                String result = entryValue.get(0);
                                if (result == null) {
                                    formDataBodyBuilder.append(entryKey).append("=").append(StringUtils.EMPTY);
                                    continue;
                                }
                                formDataBodyBuilder.append(entryKey).append("=").append(URLEncoder.encode(result, charset.name())).append("&");
                            }
                        }
                    } catch (UnsupportedEncodingException e) {
                        throw new RuntimeException(e);
                    }

                    String formDataBodyString = StringUtils.EMPTY;
                    if (formDataBodyBuilder.length() > 0) {
                        formDataBodyString = formDataBodyBuilder.substring(0, formDataBodyBuilder.length() - 1);
                    }
                    byte[] bodyBytes = formDataBodyString.getBytes(charset);
                    int contentLength = bodyBytes.length;
                    ServerHttpRequestDecorator decorator = new ServerHttpRequestDecorator(exchange.getRequest()) {

                        @Override
                        public HttpHeaders getHeaders() {
                            HttpHeaders httpHeaders = new HttpHeaders();
                            httpHeaders.putAll(super.getHeaders());
                            if (contentLength > 0) {
                                httpHeaders.setContentLength(contentLength);
                            } else {
                                httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                            }
                            return httpHeaders;
                        }

                        @Override
                        public Flux<DataBuffer> getBody() {
                            return DataBufferUtils.read(new ByteArrayResource(bodyBytes), new NettyDataBufferFactory(ByteBufAllocator.DEFAULT), contentLength);
                        }
                    };
                    ServerWebExchange mutateExchange = exchange.mutate().request(decorator).build();
                    gatewayContext.setRequestBody(formDataBodyString);
                    ServerHttpResponseDecorator decoratedResponse = this.getServerHttpResponseDecorator(exchange , gatewayContext);
                    return this.doChainFilter(mutateExchange.mutate().response(decoratedResponse).build() , chain , gatewayContext);
                }));
    }

    /**
     * 读取json格式的 request body数据
     *
     * @param exchange 交换机
     * @param chain 过滤器链
     * @return 数据对象
     */
    private Mono<Void> readBody(ServerWebExchange exchange, GatewayFilterChain chain, GatewayContext gatewayContext) {

        return DataBufferUtils.join(exchange.getRequest().getBody())
                .flatMap(dataBuffer -> {
                    DataBufferUtils.retain(dataBuffer);
                    Flux<DataBuffer> cachedFlux = Flux.defer(() -> Flux.just(dataBuffer.slice(0, dataBuffer.readableByteCount())));
                    ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(exchange.getRequest()) {

                        @Override
                        public Flux<DataBuffer> getBody() {
                            return cachedFlux;
                        }
                    };
                    ServerHttpResponseDecorator decoratedResponse = this.getServerHttpResponseDecorator(exchange , gatewayContext);
                    ServerWebExchange mutatedExchange = exchange.mutate().request(mutatedRequest).response(decoratedResponse).build();
                    return ServerRequest.create(mutatedExchange, codecConfigurer.getReaders())
                            .bodyToMono(String.class)
                            .doOnNext(gatewayContext::setRequestBody).then(this.doChainFilter(mutatedExchange , chain , gatewayContext));
                });
    }

    private ServerHttpResponseDecorator getServerHttpResponseDecorator(ServerWebExchange exchange , GatewayContext gatewayContext) {
        ServerHttpResponse originalResponse = exchange.getResponse();
        return new ServerHttpResponseDecorator(originalResponse) {

            ///private StringBuilder builder = new StringBuilder();

            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                if (body instanceof Flux) {
                    @SuppressWarnings("unchecked")
                    Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
                    /*return super.writeWith(fluxBody.map(dataBuffer -> {
                        byte[] content = new byte[dataBuffer.readableByteCount()];
                        dataBuffer.read(content);
                        //释放掉内存
                        DataBufferUtils.release(dataBuffer);
                        String responseBody = new String(content, StandardCharsets.UTF_8);
                        builder.append(responseBody);
                        gatewayContext.setResponseBody(builder.toString());
                        byte[] uppedContent = new String(content, StandardCharsets.UTF_8).getBytes();
                        return bufferFactory.wrap(uppedContent);
                    }));*/
                    return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                        DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
                        //dataBuffer合并成一个，解决获取结果不全问题
                        DataBuffer join = dataBufferFactory.join(dataBuffers);
                        byte[] content = new byte[join.readableByteCount()];
                        join.read(content);
                        String responseBody = new String(content, StandardCharsets.UTF_8);
                         /*builder.append(responseBody);
                           gatewayContext.setResponseBody(builder.toString());*/
                        gatewayContext.setResponseBody(responseBody);
                        //修改响应结果长度
                        byte[] responseBytes = responseBody.getBytes(StandardCharsets.UTF_8);
                        //originalResponse.getHeaders().setContentLength(responseBytes.length);
                        return dataBufferFactory.wrap(responseBytes);
                    }));
                }
                return super.writeWith(body);
            }
        };

    }

}